import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from datetime import datetime
import os
from typing import List, Dict, Tuple

class FactorDevelopmentVisualizer:
    """因子发展可视化工具"""
    
    def __init__(self, data_dir: str = "../../数据管理/training_data/factor_development"):
        """初始化因子发展可视化器"""
        self.data_dir = data_dir
        os.makedirs(self.data_dir, exist_ok=True)
        self.factor_data = {}  # 存储因子数据: {因子ID: 发展数据}
    
    def load_factor_data(self, factor_id: str) -> bool:
        """加载因子发展数据"""
        data_file = os.path.join(self.data_dir, f"{factor_id}_development.csv")
        if not os.path.exists(data_file):
            print(f"因子 {factor_id} 的数据文件不存在")
            return False
        
        try:
            df = pd.read_csv(data_file, parse_dates=["timestamp"])
            self.factor_data[factor_id] = df
            return True
        except Exception as e:
            print(f"加载因子数据失败: {str(e)}")
            return False
    
    def save_factor_data(self, factor_id: str, data: List[Dict]):
        """保存因子发展数据"""
        df = pd.DataFrame(data)
        if "timestamp" not in df.columns:
            df["timestamp"] = datetime.now()
        
        data_file = os.path.join(self.data_dir, f"{factor_id}_development.csv")
        # 如果文件存在则追加，否则创建
        if os.path.exists(data_file):
            existing_df = pd.read_csv(data_file, parse_dates=["timestamp"])
            df = pd.concat([existing_df, df], ignore_index=True)
        
        df.to_csv(data_file, index=False)
        self.factor_data[factor_id] = df
    
    def plot_factor_growth(self, factor_ids: List[str], metrics: List[str] = ["power", "experience"]):
        """绘制因子成长曲线"""
        fig, axes = plt.subplots(len(metrics), 1, figsize=(12, 6 * len(metrics)), sharex=True)
        if len(metrics) == 1:
            axes = [axes]
        
        for i, metric in enumerate(metrics):
            ax = axes[i]
            for factor_id in factor_ids:
                if factor_id not in self.factor_data:
                    if not self.load_factor_data(factor_id):
                        continue
                
                df = self.factor_data[factor_id]
                if metric not in df.columns:
                    print(f"因子 {factor_id} 不包含 {metric} 数据")
                    continue
                
                ax.plot(df["timestamp"], df[metric], label=factor_id)
            
            ax.set_title(f"因子{metric}成长曲线")
            ax.set_ylabel(metric)
            ax.grid(True)
            ax.legend()
        
        axes[-1].set_xlabel("时间")
        plt.tight_layout()
        return fig
    
    def plot_factor_comparison(self, factor_ids: List[str], metric: str = "power"):
        """绘制因子间对比图"""
        current_values = []
        valid_ids = []
        
        for factor_id in factor_ids:
            if factor_id not in self.factor_data:
                if not self.load_factor_data(factor_id):
                    continue
            
            df = self.factor_data[factor_id]
            if metric not in df.columns:
                print(f"因子 {factor_id} 不包含 {metric} 数据")
                continue
            
            # 获取最新值
            latest_value = df.iloc[-1][metric]
            current_values.append(latest_value)
            valid_ids.append(factor_id)
        
        if not valid_ids:
            print("没有可比较的因子数据")
            return None
        
        fig, ax = plt.subplots(figsize=(10, 6))
        bars = ax.bar(valid_ids, current_values, color=np.random.rand(len(valid_ids), 3))
        
        ax.set_title(f"因子{metric}对比")
        ax.set_ylabel(metric)
        ax.grid(axis="y")
        
        # 在柱状图上标注数值
        for bar in bars:
            height = bar.get_height()
            ax.text(bar.get_x() + bar.get_width()/2., height,
                    f'{height:.1f}', ha='center', va='bottom')
        
        plt.tight_layout()
        return fig